﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Trie
{
    public class Trie
    {
        private TrieVertex root;

        public Trie()
        {
            this.root = new TrieVertex();
            this._AddWord(this.root, string.Empty);
        }

        public void AddWord(string word)
        {
            _AddWord(root, word);
        }

        private void _AddWord(TrieVertex vertex, string word)
        {
            if (string.Empty == word)
                vertex.Words += 1;
            else
            {
                vertex.Prefixes += 1;
                int firstCharacterIndex = GetFirstCharacterIndex(word);
                if (null == vertex.Edges[firstCharacterIndex])
                    vertex.Edges[firstCharacterIndex] = new TrieVertex();
                word = RemoveLeftMostCharacter(word);
                _AddWord(vertex.Edges[firstCharacterIndex], word);
            }
        }

        public int CountWords(string word)
        {
            return _CountWords(this.root, word);
        }

        private int _CountWords(TrieVertex vertex, string word)
        {
            int firstCharacterIndex = GetFirstCharacterIndex(word);
            if (string.IsNullOrEmpty(word))
                return vertex.Words;
            else if (null == vertex.Edges[firstCharacterIndex])
                return 0;
            else
            {
                word = RemoveLeftMostCharacter(word);
                return _CountWords(vertex.Edges[firstCharacterIndex], word);
            }
        }

        public int CountWords(string word, int missingLetters)
        {
            return _CountWords(this.root, word, missingLetters);
        }

        private int _CountWords(TrieVertex vertex, string word, int missingLetters)
        {
            int firstCharacterIndex = GetFirstCharacterIndex(word);
            if (string.IsNullOrEmpty(word))
                return vertex.Words;
            else if (null == vertex.Edges[firstCharacterIndex] && 0 == missingLetters)
                return 0;
            else if (null == vertex.Edges[firstCharacterIndex])
            {
                word = RemoveLeftMostCharacter(word);
                return _CountWords(vertex.Edges[firstCharacterIndex], word, missingLetters - 1);
            }
            else
            {
                int r = _CountWords(vertex, word, missingLetters - 1);
                word = RemoveLeftMostCharacter(word);
                r += _CountWords(vertex.Edges[firstCharacterIndex], word, missingLetters);
                return r;
            }
        }

        public int CountPrefixes(string prefix)
        {
            return _CountPrefixes(this.root, prefix);
        }

        private int _CountPrefixes(TrieVertex vertex, string prefix)
        {
            int firstCharacterIndex = GetFirstCharacterIndex(prefix);
            if (string.IsNullOrEmpty(prefix))
                return vertex.Prefixes;
            else if (null == vertex.Edges[firstCharacterIndex])
                return 0;
            else
            {
                prefix = RemoveLeftMostCharacter(prefix);
                return _CountPrefixes(vertex.Edges[firstCharacterIndex], prefix);
            }
        }

        private int GetFirstCharacterIndex(string word)
        {
            if (string.IsNullOrEmpty(word))
                return -1;
            int index = ((int)(word[0])) - (int)'a';
            return index;
        }

        private string RemoveLeftMostCharacter(string word)
        {
            if (!string.IsNullOrEmpty(word))
                return word.Substring(1, word.Length-1);
            else
                return string.Empty;
        }
    }
}
